package main.java.lpr;

import java.awt.image.BufferedImage;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.bytedeco.opencv.opencv_core.Mat;

import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.util.NDImageUtils;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;

/**
 * PlateCharDetection
 * @date: 2022-02-23 10:01
 * @author: alvin
 * @copyright: Copyright © 2022 YZZH. All rights reserved.
 * @description: 基于hylperlpr端到端车牌号识别
 */
public class PlateCharDetection {
    private static String model_path = "models/lpr_gru";

    private static ZooModel<Image, PlateChar> model = null;

    private static final List<String> CLASSES = Arrays.asList("京", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏",
            "浙", "皖", "闽", "赣", "鲁",
            "豫", "鄂", "湘", "粤", "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新",
            "0", "1", "2", "3", "4", "5", "6", "7", "8", "9",
            "A", "B", "C", "D", "E", "F", "G", "H", "J", "K", "L", "M", "N", "P", "Q", "R", "S",
            "T", "U", "V", "W", "X", "Y", "Z",
            "港", "学", "使", "警", "澳", "挂", "军", "北", "南", "广", "沈", "兰", "成", "济",
            "海", "民", "航", "空");


    static {
        
        LprE2eTranslator translator = LprE2eTranslator.builder()
                .addTransform(a -> NDImageUtils.resize(a, 164, 48).transpose(1, 0, 2))
                .optClasses(CLASSES)
                .build();

        Criteria<Image, PlateChar> criteria = Criteria.builder()
                .setTypes(Image.class, PlateChar.class)
                .optModelPath(Paths.get(model_path))
                .optEngine("TensorFlow")
                .optModelName("saved_model")
                .optTranslator(translator)
                .optProgress(new ProgressBar())
                .build();

        try {
            model = criteria.loadModel();
            System.out.println(model);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }


    /**
     * 数据集对象识别 164*48的图片
     * 
     * @param img
     * @return
     */
    public static List<PlateChar> detection(Mat img, double thresh) {
        List<PlateChar> plates = new ArrayList<PlateChar>();
        if (img == null || img.empty()) {
            return plates;
        }

        // 转换成BufferedImage
        OpenCVFrameConverter.ToMat matConv = new OpenCVFrameConverter.ToMat();
        Java2DFrameConverter biConv = new Java2DFrameConverter();
        Frame frame = matConv.convert(img);
        BufferedImage bimg = biConv.getBufferedImage(frame);

        // 释放
        frame.close();
        biConv.close();
        matConv.close();

        // 转换Image
        Image djl_img = ImageFactory.getInstance().fromImage(bimg); 

        // 构建预测者
        Predictor<Image, PlateChar> predictor = model.newPredictor();

        try {
            // 检测出全部对象
            PlateChar plateChar = predictor.predict(djl_img);
            if (plateChar.getConfidence() > thresh) {
                plates.add(plateChar);
            } 
        } catch (Exception e) {
            e.printStackTrace();
        }

        // src.close();
        predictor.close();
        return plates;
    }

}
